/*
 * meituan.com Inc.
 * Copyright (c) 2010-2018 All Rights Reserved.
 */
package java.zk.example;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.List;
import java.util.Random;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * <p>
 * 主节点
 * </p>
 * @author LvJing
 * @version $Id:Master.java v1.0 2018/9/1 下午2:36 LvJing Exp $
 */
public class Master implements Watcher {

    private static final Logger log = LoggerFactory.getLogger(Master.class);

    private ZooKeeper zk;

    private String serverId = Integer.toHexString(new Random().nextInt());

    private String hostPort;

    private volatile boolean isMaster;

    private static AtomicInteger count = new AtomicInteger();

    private WorkersCache workersCache;

    public Master(String hostPort) {
        this.hostPort = hostPort;
    }

    public void startZK() throws IOException {
        this.zk = new ZooKeeper(hostPort, 15000, this);
        //bootstrap();
    }

    public void stopZK() throws InterruptedException {
        this.zk.close();
    }

    @Override
    public void process(WatchedEvent watchedEvent) {
        System.out.println(watchedEvent);
    }

    public void bootstrap() {
        createParent("/workers", "");
        createParent("/assign", "");
        createParent("/tasks", "");
        createParent("/status", "");
    }

    private void createParent(String path, String data) {
        this.zk.create(path, data.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT,
                (rc, path1, ctx, name) -> {
                    switch (KeeperException.Code.get(rc)) {
                        case CONNECTIONLOSS:
                            createParent(path, data);
                            break;
                        case OK:
                            log.info("[{}] created!", path);
                            break;
                        case NODEEXISTS:
                            log.warn("[{}] node already exist!", path);
                            break;
                        default:
                            log.error("[{}] create failed!", path,
                                    KeeperException.create(KeeperException.Code.get(rc)));
                    }
                }, null);
    }

    public void runForMaster() throws KeeperException, InterruptedException {
        while (true) {
            try {
                zk.create("/master", serverId.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
                isMaster = true;
                break;
            } catch (KeeperException.NodeExistsException e) {
                isMaster = false;
                break;
            } catch (KeeperException.ConnectionLossException e) {
                // 忽略本次请求的连接异常，可能只是本次请求失败的原因。这时可能已经在服务端创建好主节点，但是需要check一遍，看是否真的是
                // 当前进程创建的主节点。
            }
            if (checkMaster())
                break;
        }
    }

    /**
     * 异步调用，抢占主节点
     */
    public void runForMasterAsync() {
        zk.create("/master", serverId.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL,
                (rc, path, ctx, name) -> {
                    /*
                    rc:回调结果，使用int表示，异常等信息也是通过int值返回
                    path:创建节点时填写的path
                    ctx:上下文信息
                    name:创建的znode名称
                     */
                    switch (KeeperException.Code.get(rc)) {
                        case CONNECTIONLOSS:
                            // 抢占主节点，连接丢失，不用再次重试，只需检查主节点是否已经创建
                            log.warn("{} 连接丢失", serverId);
                            checkMasterAsync();
                            break;
                        case OK:
                            // 创建成功，运行中
                            isMaster = true;
                            log.info("{} 抢占成功，上位Master", serverId);
                            masterOpt();
                            break;
                        case NODEEXISTS:
                            // 主节点已存在，变为备份主节点，监视/master
                            log.info("{} 抢占Master发现节点已存在", serverId);
                            masterExists();
                            break;
                        default:
                            isMaster = false;
                            log.error("{} 抢占Master节点出现未知错误", serverId,
                                    KeeperException.create(KeeperException.Code.get(rc), path));
                    }
                }, null);
    }

    /**
     * 监视/master是否存在，存在即监视，不存在则再次抢占。监视回调，做节点删除时，抢占节点
     */
    public void masterExists() {
        zk.exists("/master", event -> {
            // /master节点删除，备份主节点开始竞争
            if (event.getType().equals(Event.EventType.NodeDeleted) && event.getPath().equals("/master")) {
                log.info("{} 备份主节点watch机制回调，Master节点宕机", serverId);
                runForMasterAsync();
            }
        }, (rc, path, ctx, stat) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    log.warn("{} 连接丢失，再次重试", serverId);
                    masterExists();
                    break;
                case NONODE:
                    log.warn("Master节点不存在");
                    runForMasterAsync();
                case OK:
                    if (stat == null) {
                        log.warn("Master节点不存在");
                        runForMasterAsync();
                    }
                    break;
                default:
                    // 发生意外情况
                    log.warn("{} 发生意外情况， {}", serverId, KeeperException.Code.get(rc));
                    checkMasterAsync();
            }
        }, null);
    }

    /**
     * 获取从节点列表信息并分配任务，同时监视从节点列表变化，一旦有变化，把对应改变的从节点重新分配任务
     */
    public void getWorkers() {
        zk.getChildren("/workers", event -> {
            if (event.getType().equals(Event.EventType.NodeChildrenChanged)) {
                if (!event.getPath().equals("/workers")) {
                    log.warn("从节点列表监视点对应的path被改变");
                    return;
                }
                getWorkers();// 再次获取从节点列表
            }
        }, (rc, path, ctx, children) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    log.warn("{} 连接丢失，重试", serverId);
                    getWorkers();
                    break;
                case OK:
                    log.info("成功获取从节点列表，size:{}", children.size());
                    // 重新为从节点列表分配任务
                    reassignAndSet(children);
                    break;
                default:
                    log.error("{} 获取从节点列表失败", serverId, KeeperException.create(KeeperException.Code.get(rc)));
            }
        }, null);
    }

    /**
     * 主节点监视未分配任务变化，即/tasks。一旦有节点添加在/tasks下，则为其分配一个从节点，
     * 在/assign下创建分配记录，路径规则为/assign/worker-id/task-id。
     * 分配记录创建过后，删除/tasks下相应节点，因为下面的节点都是没有分配的任务。
     */
    public void getTasks() {
        zk.getChildren("/tasks", event -> {
            if (event.getType().equals(Event.EventType.NodeChildrenChanged)) {
                if (!event.getPath().equals("/tasks")) {
                    log.warn("/tasks 路径名发生改变");
                    return;
                }
                getTasks();
            }
        }, (rc, path, ctx, children) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    log.warn("{} 连接丢失，再次尝试getChildren", serverId);
                    getTasks();
                    break;
                case OK:
                    if (children != null && children.size() > 0) {
                        log.info("新增待分配任务{}个", children.size());
                        assignTasks(children);
                    }
                    break;
                default:
                    log.error("{} getChildren，出现未知错误", serverId, KeeperException.create(KeeperException.Code.get(rc)));
            }
        }, null);
    }

    /**
     * 分配任务。在从节点列表随机选出一个从节点，分配给该任务，在/assign下创建/assign/worker-id/task-id，
     * 同时从/tasks下删除该任务，表示该任务已分配给从节点执行。
     */
    public void assignTasks(List<String> tasks) {
        for (String task : tasks) {
            assignTask(task);
        }
    }

    private void assignTask(String task) {
        zk.getData("/tasks/" + task, false, (rc, path, ctx, data, stat) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    log.warn("{} 连接丢失，重试", serverId);
                    assignTask((String) ctx);
                    break;
                case OK:
                    int randomId = new Random().nextInt(workersCache.workers.size());
                    String designatedWorker = workersCache.workers.get(randomId);
                    log.info("{} 任务分配给了 {}", ctx, designatedWorker);
                    createAssignPath(designatedWorker, (String) ctx, data);
            }
        }, task);
    }

    private void createAssignPath(String designatedWorker, String taskId, byte[] taskData) {
        String assignmentPath = "/assign/" + designatedWorker + "/" + taskId;
        zk.create(assignmentPath, taskData, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT,
                (rc, path, ctx, name) -> {
                    switch (KeeperException.Code.get(rc)) {
                        case CONNECTIONLOSS:
                            log.warn("{} 连接丢失，重试尝试创建分配路径", serverId);
                            createAssignPath(designatedWorker, taskId, (byte[]) ctx);
                            break;
                        case OK:
                            log.info("{} 分配路径创建成功", path);
                            deleteTask(taskId);
                            break;
                        case NODEEXISTS:
                            log.warn("{} 分配路径已经存在，无法分配", path);
                            break;
                        default:
                            log.error("{} 分配路径创建时出现未知错误", path, KeeperException.create(KeeperException.Code.get(rc)));
                    }
                }, taskData);
    }

    private void deleteTask(String task) {
        zk.delete("/tasks/" + task, 0, (rc, path, ctx) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    log.warn("{} 连接丢失，重试尝试删除任务", serverId);
                    deleteTask((String) ctx);
                    break;
                case OK:
                    log.info("{} 任务删除成功", task);
                    break;
                default:
                    log.error("{} 任务删除时出现未知错误", task, KeeperException.create(KeeperException.Code.get(rc)));
            }
        }, task);
    }

    /**
     * 重新为从节点分配任务并设置
     */
    private void reassignAndSet(List<String> children) {
        List<String> deletedWorkers;
        if (workersCache == null) {
            // 初始化workersCache
            workersCache = new WorkersCache(children);
            deletedWorkers = null;// 第一次获取从节点列表，无需做任何事情
        } else {
            log.info("获取已被删除的节点");
            deletedWorkers = workersCache.removeAndSet(children);
        }
        if (deletedWorkers != null && deletedWorkers.size() > 0) {
            for (String deletedWorker : deletedWorkers) {
                // 获取已删除的节点任务，重新分配
                // TODO: 2018/9/4 这里涉及查询已删除节点的任务状态为【未成功】，需要再次重新分配从节点去执行
            }
        }
    }

    private void masterOpt() {
        log.info("{} 作为Master节点运行中!", serverId);
        getWorkers();
        log.info("{} 监视从节点列表变化", serverId);
        getTasks();
        log.info("{} 监视新任务并分配给从节点", serverId);
    }

    /**
     * 异步调用，检查是否创建了主节点
     */
    private void checkMasterAsync() {
        zk.getData("/master", false, (rc, path, ctx, data, stat) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    log.warn("{} 连接丢失，再次重试", serverId);
                    checkMasterAsync();
                    break;
                case NONODE:
                    runForMasterAsync();
                    break;
                case OK:
                    isMaster = new String(data).equals(serverId);
                    // 未抢占到master，监视master
                    if (!isMaster)
                        masterExists();
                    break;
                default:
                    // 发生未知错误
                    log.error("{} 获取master数据错误", serverId, KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }, null);
    }

    private boolean checkMaster() throws KeeperException, InterruptedException {
        while (true) {
            try {
                // 获取主节点数据，检查是否为
                byte[] data = zk.getData("/master", false, new Stat());
                isMaster = new String(data).equals(serverId);
                return true;
            } catch (KeeperException.NoNodeException e) {
                return false;// 主节点不存在，再次创建主节点
            } catch (KeeperException.ConnectionLossException e) {
                // 再次尝试请求连接。
            }
        }
    }

    public static void main(String[] args) {
        ThreadPoolExecutor executor = new ThreadPoolExecutor(0, 10, 0, TimeUnit.MILLISECONDS,
                new LinkedBlockingDeque<>());
        count.set(10);
        for (int i = 0; i < 10; i++) {
            executor.execute(() -> {
                try {
                    Master master = new Master("127.0.0.1:2181");
                    master.startZK();
                    master.runForMasterAsync();
                    while (!master.isMaster) {
                    }
                    Thread.sleep(5000);
                    log.info("{} 即将宕机，ZK会删除master节点", master.serverId);
                    log.info("剩余线程数量:{}", count.decrementAndGet());
                    master.stopZK();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            });
        }
    }

}
